/*
 * Description: mount host shared folder
 */

#include "share_folder_helper.h"

#include <future>
#include <sstream>
#include <windows.h>

#include <WtsApi32.h>
#include "winnetwk.h"

#include "../common/guest_utils.h"
#include "../common/text_utils.h"
#include "../common/session_utils.h"
#include "../common/win_update_utils.h"
#include "log.h"

#pragma comment(lib, "Mpr.lib")

using namespace COMMON;

namespace {
const int NETRESOURCE_BUFF_SIZE = 16384;
const std::string MOUNT_POINT_ROOT_REG = "\\Software\\Microsoft\\Windows\\CurrentVersion\\Explorer\\MountPoints2\\";
const std::string REMOTE_PREFIX = "\\\\HwFs\\";
}  // namespace

namespace GUEST {
void ShareFolderHelper::Init(std::shared_ptr<WINCONTROL::SendToHost> sendtoHost)
{
    std::unique_lock<std::mutex> lk(mountMutex);
    if (isInit) {
        LOG_INFO("ShareFolderHelper has init.");
        return;
    }
    this->sendtoHost = sendtoHost;
    this->isInit = true;
    LOG_INFO("ShareFolderHelper init.");
    SessionUtils::RegSessionChangeCallback(sessionCallbackName_,
        std::bind(&ShareFolderHelper::OnWindowsSessionChange, &SHARE_FOLDER_HELPER, std::placeholders::_1));
}

void ShareFolderHelper::AddSharedFolder(const std::string &shareName)
{
    std::unique_lock<std::mutex> lk(mountMutex);
    std::string shareNameGBK = TextUtils::UTF8ToGBK(shareName);

    auto mountedList = GetAllMountedList();
    for (auto &item : mountedList) {
        if (item.first == shareNameGBK) {
            return;
        }
    }

    if (!SessionUtils::IsUserLogon()) {
        LOG_WARN("The user has not logged in.");
        return;
    }

    auto letter = GuestUtils::GetLastAvailableLetter();
    std::string remote = REMOTE_PREFIX + shareNameGBK;
    if (letter.length() == 0) {
        LOG_ERR("available letter is not exist!");
        return;
    }
    NETRESOURCEA *nr = new NETRESOURCEA;
    nr->dwType = RESOURCETYPE_DISK;
    nr->lpLocalName = (LPSTR)letter.c_str();
    nr->lpRemoteName = (LPSTR)remote.c_str();
    nr->lpProvider = NULL;
    auto ret = WNetAddConnection2A(nr, NULL, NULL, CONNECT_TEMPORARY | CONNECT_REDIRECT);
    LOG_WARN("mount shared folder ret: {}", ret);
    UpdateMountPointRegistry(shareNameGBK);
    delete nr;
    nr = nullptr;
}

void ShareFolderHelper::RemoveSharedFolder(const std::string &shareName)
{
    std::unique_lock<std::mutex> lk(mountMutex);
    std::string shareNameGBK = TextUtils::UTF8ToGBK(shareName);
    MountItem target;
    auto mountedList = GetAllMountedList();
    for (auto &item : mountedList) {
        if (item.first == shareNameGBK) {
            target = item;
            break;
        }
    }
    DWORD ret = WNetCancelConnection2A(target.second.c_str(), CONNECT_UPDATE_PROFILE, TRUE);
    LOG_WARN("remove shared folder ret: {}", ret);
}

int ShareFolderHelper::Deinit()
{
    this->sendtoHost = nullptr;
    this->isInit = false;
    SessionUtils::UnRegSessionChangeCallback(sessionCallbackName_);
    return 0;
}

void ShareFolderHelper::UpdateMountPointRegistry(const std::string &shareName)
{
    std::stringstream mountPointKey0;
    mountPointKey0 << MOUNT_POINT_ROOT_REG << "##HwFs#" << shareName;
    GuestUtils::SetRegValueForCurrentUser(mountPointKey0.str(), "_LabelFromReg", REG_SZ, shareName);

    std::stringstream mountPointKey1;
    mountPointKey1 << MOUNT_POINT_ROOT_REG << "HwFs#" << shareName;

    GuestUtils::SetRegValueForCurrentUser(mountPointKey1.str(), "_LabelFromReg", REG_SZ, shareName);
}

void ShareFolderHelper::UpdateMountPointRegistry()
{
    std::stringstream cmd;
    for (auto &item : GetAllMountedList()) {
        UpdateMountPointRegistry(item.first);
    }
}

void ShareFolderHelper::OnWindowsSessionChange(int32_t state)
{
    LOG_INFO("guest session change: {}.", state);
    if (WTS_SESSION_LOGON == state || WTS_SESSION_UNLOCK == state) {
        LOG_INFO("guest session change, update net driver info.");
        UpdateMountPointRegistry();
    }
    // upgrade agent
    if (WTS_SESSION_LOGON == state) {
        LOG_INFO("guest session change, logon upgrade agent.");
        WinUpdateUtils::UpgradeAgent();
    }
    // notify host
    if (this->sendtoHost != nullptr) {
        this->sendtoHost->GuestSessionChange(state);
    }
}

std::vector<MountItem> ShareFolderHelper::GetAllMountedList()
{
    std::vector<MountItem> result;

    HANDLE handle = NULL;
    DWORD ret = WNetOpenEnumW(RESOURCE_CONNECTED, RESOURCETYPE_ANY, 0, NULL, &handle);
    if (ret != NO_ERROR) {
        LOG_ERR("WNetOpenEnum failed with error {} ", ret);
        return result;
    }
    DWORD buffSize = NETRESOURCE_BUFF_SIZE;
    DWORD maxSize = NETRESOURCE_BUFF_SIZE / sizeof(NETRESOURCEA);
    LPNETRESOURCEA resource;

    resource = (LPNETRESOURCEA)GlobalAlloc(GPTR, buffSize);
    if (resource == NULL) {
        LOG_ERR("Failed to allocate memory");
        WNetCloseEnum(handle);
        return result;
    }
    do {
        DWORD count = -1;
        ret = WNetEnumResourceA(handle, &count, resource, &buffSize);
        if (ret == ERROR_MORE_DATA) {
            GlobalFree(resource);
            resource = (LPNETRESOURCEA)GlobalAlloc(GPTR, buffSize);
            if (resource == NULL) {
                LOG_ERR("Failed to allocate memory");
                break;
            }
        }
        if (ret != NO_ERROR || count > maxSize) {
            continue;
        }
        for (int i = 0; i < count; i++) {
            if (resource[i].lpLocalName == nullptr) {
                continue;
            }
            std::string letter(resource[i].lpLocalName);
            std::string remote(resource[i].lpRemoteName);
            if (remote.find(REMOTE_PREFIX) != 0) {
                continue;
            }
            result.emplace_back(remote.substr(REMOTE_PREFIX.length()), letter);
        }
    } while (ret != ERROR_NO_MORE_ITEMS);

    GlobalFree(resource);
    WNetCloseEnum(handle);
    return result;
}

void ShareFolderHelper::RemoveAllSharedFolder()
{
    auto mountedList = GetAllMountedList();
    for (auto &item : mountedList) {
        WNetCancelConnection2A(item.second.c_str(), CONNECT_UPDATE_PROFILE, TRUE);
    }
}

}  // namespace GUEST